import os
import sys
import json
import cv2
import numpy as np
import label
import tensorflow as tf

# 设置训练集、测试集的地址
base_dir = 'D:\\deep_learning\\model'
saved_files = 'D:\\deep_learning\\model\\output_files'

# 创建模型结果文件
if not os.path.exists(saved_files):
    os.makedirs(saved_files)

# 获取训练集目录及种类数
labels = label.getFruitLabel()
num_classes = len(labels)


def test_model(name="", filePath=""):
    model_out_dir = os.path.join(saved_files, f'{name}.h5')
    print(model_out_dir)
    if not os.path.exists(model_out_dir):
        print("No saved model found")
        exit(0)
    # 读取模型
    model = tf.keras.models.load_model(model_out_dir)

    # 获取图片并处理
    image = cv2.imread(filePath)
    image = cv2.resize(image, (100, 100))

    # 设置图片维度
    data = np.ndarray(shape=(1, 100, 100, 3), dtype=np.int_)
    image_array = np.asarray(image)
    data[0] = image_array

    # 预测结果
    result = model.predict(data, 1)
    fruitClass = labels[result.argmax(axis=-1)[0]]
    return label.getFruitDict(fruitClass)


def main():
    filePath = sys.argv[1]
    result = test_model(name='test', filePath=filePath)
    print(json.dumps(result))
    sys.stdout.flush()


if __name__ == '__main__':
    main()
